import os
import csv

training_class_dir = './Final_Training/Images/'
training_label_output_dir = './Final_Training_Label/'
training_label_csv_filename = 'label.csv'

label_list = []

num_class = 0

class_dir_all = os.listdir(training_class_dir)
class_dir_all.sort()
print(class_dir_all)

for class_dir in class_dir_all:
    # Get the full path
    training_label_dir = training_class_dir + class_dir
    print(class_dir)
    # Get the csv file
    csv_filename = training_label_dir + '/GT-' + class_dir + '.csv'
    with open(csv_filename) as fi:
        reader = csv.reader(fi)

        # Get each img's label
        for row in reader:
            if reader.line_num < 2:
                continue

            img_filename, width, height, x1, x2, y1, y2, label = row[0].split(';')

            if len(label) == 1:
                label_for_name = '0' + label
            else:
                label_for_name = label
            img_filename = label_for_name + '_' + img_filename.split('.')[0] + '.jpg'

            label_list.append([img_filename, label])

    # We only get 43 classes
    num_class += 1
    if num_class == 43:
        break
# Write the label_list to label csv file

with open(training_label_output_dir + training_label_csv_filename, 'w') as fw:
    writer = csv.writer(fw)
    for row in label_list:
        writer.writerow(row)

print('pre-dataset is created.')
